import matplotlib.pyplot as plt
import numpy as np
from PIL import Image


def svd_compression(image_path, k):
    # 读取图像并转换为灰度
    img = Image.open(image_path).convert("L")
    img_array = np.array(img)

    # 进行SVD分解
    U, s, VT = np.linalg.svd(img_array, full_matrices=False)

    # 选取前k个奇异值
    S = np.diag(s[:k])
    U_k = U[:, :k]
    VT_k = VT[:k, :]

    # 压缩图像
    # 使用奇异值分解（SVD）的结果对图像进行压缩
    # U_k: 左奇异向量矩阵的截断版本
    # S: 奇异值矩阵的截断版本
    # VT_k: 右奇异向量矩阵的转置的截断版本
    # np.dot() 函数用于矩阵乘法，通过乘以这些截断矩阵来实现图像的压缩
    compressed_img = np.dot(U_k, np.dot(S, VT_k))


    return img_array, compressed_img


# 图像路径和奇异值数量
image_path = 'images/svd_test_image.png'  # 替换为你的图像路径
k_values = [5, 10, 50, 100]  # 不同的奇异值数量

# 显示图像和压缩效果
original_img, compressed_img = svd_compression(image_path, k=50)

plt.figure(figsize=(12, 6))

# 显示原始图像
plt.subplot(1, 2, 1)
plt.title("Original Image")
plt.imshow(original_img, cmap='gray')
plt.axis('off')

# 显示压缩图像
plt.subplot(1, 2, 2)
plt.title("Compressed Image with k=50")
plt.imshow(compressed_img, cmap='gray')
plt.axis('off')

plt.show()
